import cv2
import numpy as np
import matplotlib.pyplot as plt

def linear_transform(img, low_in=0, high_in=1, low_out=0, high_out=1):
    assert high_in >= 0 and high_in <= 1 and low_in >= 0 and low_in <= 1
    assert high_out >= 0 and high_out <= 1 and low_out >= 0 and low_out <= 1
    img = img / 255.0
    out = np.zeros_like(img)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            if img[i,j] < low_in:
                out[i,j] = low_out
            elif img[i,j] > high_in:
                out[i,j] = high_out
            else:
                k = (high_out-low_out)/(high_in-low_in)
                b = (high_in*low_out-low_in*high_out)/(high_in-low_in)
                out[i,j] = img[i,j] * k + b
    out = out * 255.0
    return k,b,out

def main():
    grayimg = cv2.imread('camera.png', cv2.IMREAD_GRAYSCALE)
    plt.figure(figsize=(8, 8))                          # 显示原始图像
    plt.subplot(2, 2, 1), plt.imshow(grayimg, cmap='gray'), plt.title('original img')
    H = cv2.equalizeHist(grayimg)                     # 直方图均衡化并显示图像
    plt.subplot(2, 2, 2), plt.imshow(H, cmap='gray'), plt.title('after histeq')
    # cv2.imwrite('H.png', H)
    _, _, I1 = linear_transform(grayimg, 0, 1, 1, 0)         # 线性变换1并显示图像
    plt.subplot(2, 2, 3), plt.imshow(I1, cmap='gray'), plt.title('after linear transform1')
    _, _, I2 = linear_transform(grayimg, 0.3, 0.7, 0.1, 0.9)   # 线性变换2并显示图像
    plt.subplot(2, 2, 4), plt.imshow(I2, cmap='gray'), plt.title('after linear transform2')
    plt.figure(num=2, figsize=(8, 8))                    # 显示各图像直方图
    plt.subplot(2, 2, 1), plt.hist(grayimg.flatten(), bins=256), plt.title('original hist')
    plt.subplot(2, 2, 2), plt.hist(H.flatten(), bins=256), plt.title('after histeq')
    plt.subplot(2, 2, 3), plt.hist(I1.flatten(), bins=256), plt.title('after linear transform1')
    plt.subplot(2, 2, 4), plt.hist(I2.flatten(), bins=256), plt.title('after linear transform2')
    plt.show()
    # 另一种画直方图的方法
    # his_ori = cv2.calcHist([grayimg], [0], None, [256], [0, 255])
    # plt.subplot(2, 2, 1), plt.plot(his_ori), plt.title('original hist')
if __name__ == '__main__':
    main()